import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
from itertools import zip_longest
import time

# 设置 matplotlib 支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 生成从 1 到 100 的数组
arr = np.arange(1, 101)
np.random.shuffle(arr)  # 打乱数组顺序

# 冒泡排序生成器函数
def bubble_sort_generator(arr):
    start_time = time.time()
    n = len(arr)
    for i in range(n):
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
            yield arr
    end_time = time.time()
    yield arr, end_time - start_time

# 选择排序生成器函数
def selection_sort_generator(arr):
    start_time = time.time()
    n = len(arr)
    for i in range(n):
        min_idx = i
        for j in range(i + 1, n):
            if arr[j] < arr[min_idx]:
                min_idx = j
        arr[i], arr[min_idx] = arr[min_idx], arr[i]
        yield arr
    end_time = time.time()
    yield arr, end_time - start_time

# 插入排序生成器函数
def insertion_sort_generator(arr):
    start_time = time.time()
    for i in range(1, len(arr)):
        key = arr[i]
        j = i - 1
        while j >= 0 and key < arr[j]:
            arr[j + 1] = arr[j]
            j -= 1
        arr[j + 1] = key
        yield arr
    end_time = time.time()
    yield arr, end_time - start_time

# 希尔排序生成器函数
def shell_sort_generator(arr):
    start_time = time.time()
    n = len(arr)
    gap = n // 2
    while gap > 0:
        for i in range(gap, n):
            temp = arr[i]
            j = i
            while j >= gap and arr[j - gap] > temp:
                arr[j] = arr[j - gap]
                j -= gap
            arr[j] = temp
            yield arr
        gap //= 2
    end_time = time.time()
    yield arr, end_time - start_time

# 归并排序辅助函数，合并两个有序子数组
def merge(arr, l, m, r):
    n1 = m - l + 1
    n2 = r - m
    L = [0] * n1
    R = [0] * n2
    for i in range(0, n1):
        L[i] = arr[l + i]
    for j in range(0, n2):
        R[j] = arr[m + 1 + j]
    i = 0
    j = 0
    k = l
    while i < n1 and j < n2:
        if L[i] <= R[j]:
            arr[k] = L[i]
            i += 1
        else:
            arr[k] = R[j]
            j += 1
        k += 1
        yield arr
    while i < n1:
        arr[k] = L[i]
        i += 1
        k += 1
        yield arr
    while j < n2:
        arr[k] = R[j]
        j += 1
        k += 1
        yield arr

# 归并排序生成器函数
def merge_sort_generator(arr, l=0, r=None):
    start_time = time.time()
    if r is None:
        r = len(arr) - 1
    if l < r:
        m = (l + r) // 2
        yield from merge_sort_generator(arr, l, m)
        yield from merge_sort_generator(arr, m + 1, r)
        yield from merge(arr, l, m, r)
    end_time = time.time()
    yield arr, end_time - start_time

# 快速排序辅助函数，分区操作
def partition(arr, low, high):
    pivot = arr[high]
    i = low - 1
    for j in range(low, high):
        if arr[j] <= pivot:
            i = i + 1
            arr[i], arr[j] = arr[j], arr[i]
            yield arr
    arr[i + 1], arr[high] = arr[high], arr[i + 1]
    yield arr
    return i + 1

# 快速排序生成器函数
def quick_sort_generator(arr, low=0, high=None):
    start_time = time.time()
    if high is None:
        high = len(arr) - 1
    if low < high:
        pi = None
        pi_gen = partition(arr, low, high)
        try:
            while True:
                state = next(pi_gen)
                yield state
        except StopIteration as e:
            pi = e.value  # 获取分区函数的返回值
        if pi is not None:
            yield from quick_sort_generator(arr, low, pi - 1)
            yield from quick_sort_generator(arr, pi + 1, high)
    end_time = time.time()
    yield arr, end_time - start_time

# 创建图形和坐标轴，按两行三列布局
fig, axes = plt.subplots(2, 3, figsize=(15, 10))

# 设置图表标题
titles = ['冒泡排序可视化', '选择排序可视化', '插入排序可视化', '希尔排序可视化', '归并排序可视化', '快速排序可视化']
bar_rects = []
texts = []
for i, ax in enumerate(axes.flatten()):
    ax.set_title(titles[i])
    bar_rects.append(ax.bar(range(len(arr)), arr, align="edge"))
    # 设置文字大小为 14，颜色为绿色
    texts.append(ax.text(0.5, 0.9, '', transform=ax.transAxes, ha='center', fontsize=14, color='green'))

# 记录每个排序生成器是否已经结束
sort_finished = [False] * 6

# 更新函数，用于更新每一帧的图形
def update(frame):
    for i, (rects, text) in enumerate(zip(bar_rects, texts)):
        if frame[i] is not None:
            if isinstance(frame[i], tuple) and not sort_finished[i]:
                arr_state, elapsed_time = frame[i]
                for rect, val in zip(rects, arr_state):
                    rect.set_height(val)
                text.set_text(f"耗时: {elapsed_time:.4f} 秒")
                sort_finished[i] = True
            else:
                for rect, val in zip(rects, frame[i] if not isinstance(frame[i], tuple) else frame[i][0]):
                    rect.set_height(val)
    return [rect for rects in bar_rects for rect in rects] + texts

# 创建动画
sort_generators = [
    bubble_sort_generator(arr.copy()),
    selection_sort_generator(arr.copy()),
    insertion_sort_generator(arr.copy()),
    shell_sort_generator(arr.copy()),
    merge_sort_generator(arr.copy()),
    quick_sort_generator(arr.copy())
]
frames = zip_longest(*sort_generators)

ani = animation.FuncAnimation(
    fig,
    func=update,
    frames=frames,
    interval=10,
    repeat=False,
    save_count=len(arr) ** 2
)

# 显示图形
plt.show()
